#ifndef ELEMENT_PROXY_H
#define ELEMENT_PROXY_H

#include "common.h"

namespace tinyndarray {

// 前置声明
template <typename T>
class ElementProxyBase;

template <typename T>
std::ostream& operator<<(std::ostream& os, const ElementProxyBase<T>& proxy);

template <typename T>
class ElementProxyBase {
private:
    T* ptr_;       // 指向元素的指针
    Dtype dtype_;  // 元素的数据类型

public:
    ElementProxyBase(T* ptr, Dtype dtype);
    // 添加访问私有成员的公共方法
    T* ptr() const {
        return ptr_;
    }
    Dtype dtype() const {
        return dtype_;
    }

    // 类型转换运算符
    operator float32_t() const;
    operator float64_t() const;
    operator int() const;
    operator bool() const;
    void set_ptr(T* ptr);
    void set_dtype(Dtype dtype);

    // 修改为模板友元声明
    friend std::ostream& operator<< <T>(std::ostream& os, const ElementProxyBase<T>& proxy);
};

template <typename T>
ElementProxyBase<T>::ElementProxyBase(T* ptr, Dtype dtype) : ptr_{ptr}, dtype_{dtype} {}

template <typename T>
void ElementProxyBase<T>::set_ptr(T* ptr) {
    ptr_ = ptr;
}

template <typename T>
void ElementProxyBase<T>::set_dtype(Dtype dtype) {
    dtype_ = dtype;
}

template <typename T>
ElementProxyBase<T>::operator float32_t() const {
    if (dtype().tag != FP32) throw std::runtime_error("Type mismatch: not a float");
    return *static_cast<const float32_t*>(ptr());
}

template <typename T>
ElementProxyBase<T>::operator int() const {
    if (dtype().tag != INT32) throw std::runtime_error("Type mismatch: not an int");
    return *static_cast<const int32_t*>(ptr());
}

template <typename T>
ElementProxyBase<T>::operator float64_t() const {
    if (dtype().tag != FP64) throw std::runtime_error("Type mismatch: not a double");
    return *static_cast<const float64_t*>(ptr());
}

template <typename T>
ElementProxyBase<T>::operator bool() const {
    if (dtype().tag != BOOL) throw std::runtime_error("Type mismatch: not a bool");
    return *static_cast<const bool*>(ptr());
}
// 常量元素代理类
class ConstElementProxy : public ElementProxyBase<const void> {
public:
    // 添加带参数的构造函数
    ConstElementProxy(const void* ptr, Dtype dtype) : ElementProxyBase<const void>(const_cast<void*>(ptr), dtype) {}
};

// 可变元素代理类
class ElementProxy : public ElementProxyBase<void> {
private:
    // 应用一元操作的模板方法
    template <typename Op>
    ElementProxy apply_unary_operation(Op op) const {
        // 使用公共访问方法替代直接访问私有成员
        switch (dtype().tag) {
            case FP32: {
                float32_t val = *static_cast<float32_t*>(ptr());
                // 使用两个参数的构造函数
                return ElementProxy(new float32_t(op(val)), dtype());
            }
            case INT32: {
                int32_t val = *static_cast<int32_t*>(ptr());
                return ElementProxy(new int32_t(op(val)), dtype());
            }
            case FP64: {
                float64_t val = *static_cast<float64_t*>(ptr());
                return ElementProxy(new float64_t(op(val)), dtype());
            }
            case BOOL: {
                bool val = *static_cast<bool*>(ptr());
                return ElementProxy(new bool(op(val)), dtype());
            }
            default: throw std::runtime_error("Unsupported data type for operation");
        }
    }

public:
    ElementProxy(void* ptr, Dtype dtype) : ElementProxyBase<void>(ptr, dtype) {}

    void reset(void* ptr, Dtype dtype);

    ElementProxy& operator=(const ElementProxy& other);
    // 赋值运算符
    ElementProxy& operator=(float32_t value);
    ElementProxy& operator=(int32_t value);
    ElementProxy& operator=(float64_t value);
    ElementProxy& operator=(bool value);
    // 通用赋值模板
    template <typename T>
    ElementProxy& operator=(T value) {
        // 使用公共访问方法
        switch (dtype().tag) {
            case FP32: *static_cast<float*>(ptr()) = static_cast<float>(value); break;
            case INT32: *static_cast<int*>(ptr()) = static_cast<int>(value); break;
            case FP64: *static_cast<double*>(ptr()) = static_cast<double>(value); break;
            case BOOL: *static_cast<bool*>(ptr()) = static_cast<bool>(value); break;
            default: throw std::runtime_error("Unknown data type for assignment");
        }
        return *this;
    }

    // 算术运算符重载
    ElementProxy operator*(int factor) const {
        return apply_unary_operation([factor](auto a) { return a * factor; });
    }
};

// 显式实例化
template class ElementProxyBase<const void>;
template class ElementProxyBase<void>;

// 实现输出运算符
template <typename T>
std::ostream& operator<<(std::ostream& os, const ElementProxyBase<T>& proxy) {
    // 使用公共访问方法
    switch (proxy.dtype().tag) {
        case FP32: os << *static_cast<const float*>(proxy.ptr()); break;
        case FP64: os << *static_cast<const double*>(proxy.ptr()); break;
        case QINT8: os << static_cast<int32_t>(*static_cast<const int8_t*>(proxy.ptr())); break;
        case INT32: os << *static_cast<const int32_t*>(proxy.ptr()); break;
        case BOOL: os << *static_cast<const bool*>(proxy.ptr()) ? "true" : "false"; break;
        default: throw std::runtime_error("Unknown data type for output");
    }
    return os;
}

}  // namespace tinyndarray
#endif  // ELEMENT_PROXY_H